/*
 * Copyright 2013 Evan Jones.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.keyczar;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

import java.io.UnsupportedEncodingException;

import javax.crypto.spec.PBEKeySpec;

import junit.framework.TestCase;

import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.keyczar.KeyczarPBEReader.PBEKeyczarKey;
import org.keyczar.enums.KeyPurpose;
import org.keyczar.enums.KeyStatus;
import org.keyczar.exceptions.KeyczarException;
import org.keyczar.interfaces.KeyczarReader;
import org.keyczar.util.Base64Coder;
import org.keyczar.util.Util;

public class KeyczarPBEReaderTest extends TestCase {
  private static final String TEST_DATA = "./testdata";

  // Generated using C++ Keyczar
  private static final String JSON_KEY = "{\n" +
      "\"cipher\": \"AES128\"," +
      "\"hmac\": \"HMAC_SHA1\"," +
      "\"iterationCount\": 4096," +
      "\"iv\": \"z3BdMSyqfrh-qmv1YLXvFg\"," +
      "\"key\": \"ZoavQvg_IRDGG57NQIn4kjuBRyRsX3p6JPWX1jnNUBtUQyAHTd381CzKyqOZIIVt8nIkkzdN3JjtTyYMQSEE9ZTVZ_RVC1enTzLZuEL5gZCbmJyRzX1eBdpTN1bFbIt3aOMiFxjFzP-O67ErGnUHpBHmmCxmau-MBUpCd6Su-eum3SIERsaMzsMDEcivCrd5SW20HokMWCxu_GImVxyQuA\"," +
      "\"salt\": \"EwDMYR65XcUmvggvoW1GMw\"" +
      "}";

  private final static String METADATA = "{\n" +
      "   \"encrypted\": true,\n" +
      "   \"name\": \"Test\",\n" +
      "   \"purpose\": \"DECRYPT_AND_ENCRYPT\",\n" +
      "   \"type\": \"AES\",\n" +
      "   \"versions\": [ {\n" +
      "      \"exportable\": false,\n" +
      "      \"status\": \"PRIMARY\",\n" +
      "      \"versionNumber\": 1\n" +
      "   } ]\n" +
      "}\n";

  private static final String PASSPHRASE = "pass";

  @Rule
  public TemporaryFolder tempFolder = new TemporaryFolder();

  private static class EncryptedTestKeyReader implements KeyczarReader {
    @Override
    public String getKey() throws KeyczarException {
      return JSON_KEY;
    }

    @Override
    public String getKey(int version) throws KeyczarException {
      assert version == 1;
      return JSON_KEY;
    }

    @Override
    public String getMetadata() throws KeyczarException {
      return METADATA;
    }
  };

  @Test
  public void testKeySpecFromJson() throws KeyczarException {
    KeyczarPBEReader.PBEKeyczarKey pbeKey = KeyczarPBEReader.parsePBEMetadata(JSON_KEY);

    PBEKeySpec spec = KeyczarPBEReader.keySpecFromJson(pbeKey, PASSPHRASE);
    assertEquals(16*8, spec.getKeyLength());
    assertEquals(4096, spec.getIterationCount());
    final byte[] salt = {19, 0, -52, 97, 30, -71, 93, -59, 38, -66, 8, 47, -95, 109, 70, 51};
    assertArrayEquals(salt, spec.getSalt());
  }

  @Test
  public void testParseBadKeyMetadata() throws KeyczarException {
    KeyczarPBEReader.PBEKeyczarKey pbeKey = new KeyczarPBEReader.PBEKeyczarKey();
    // all null fields
    assertParseException(pbeKey, "Invalid password-based key metadata");
    pbeKey.cipher = "";
    assertParseException(pbeKey, "Invalid password-based key metadata");
    pbeKey.hmac = "";
    assertParseException(pbeKey, "Invalid password-based key metadata");
    pbeKey.iterationCount = -1;
    assertParseException(pbeKey, "Unsupported password-based encryption algorithm");

    pbeKey.cipher = "AES128";
    assertParseException(pbeKey, "Unsupported key derivation hash function");
    pbeKey.hmac = "HMAC_SHA1";
    assertParseException(pbeKey, "Invalid iteration count");
    pbeKey.iterationCount = 1;

    // this is actually "correct" metadata
    String json = Util.gson().toJson(pbeKey);
    assertNotNull(KeyczarPBEReader.parsePBEMetadata(json));
  }

  private void assertParseException(KeyczarPBEReader.PBEKeyczarKey pbeKey, String expectedMessage) {
    String json = Util.gson().toJson(pbeKey);
    try {
      KeyczarPBEReader.parsePBEMetadata(json);
      fail("expected exception");
    } catch (KeyczarException e) {
      assertEquals(expectedMessage, e.getMessage());
    }
  }

  private static byte[] convertHex(String hex) {
    String[] parts = hex.split(" ");
    byte[] out = new byte[parts.length];
    for (int i = 0; i < parts.length; ++i) {
      out[i] = (byte) Integer.parseInt(parts[i], 16);
    }
    return out;
  }

  @Test
  public void testPBKDF2() throws UnsupportedEncodingException {
    // Test vectors from http://tools.ietf.org/html/rfc6070
    assertKey(1, "0c 60 c8 0f 96 1f 0e 71 f3 a9 b5 24 af 60 12 06 2f e0 37 a6");
    assertKey(2, "ea 6c 01 4d c7 2d 6f 8c cd 1e d9 2a ce 1d 41 f0 d8 de 89 57");
    assertKey(4096, "4b 00 79 01 b7 65 48 9a be ad 49 d9 26 f7 21 d0 65 a4 29 c1");
    // This runs too slow; I wasn't patient enough to verify this :)
    // assertKey(16777216, "ee fe 3d 61 cd 4d a4 e4 e9 94 5b 3d 6b a2 15 8c 26 34 e9 84");

    String password = "passwordPASSWORDpassword";
    String salt = "saltSALTsaltSALTsaltSALTsaltSALTsalt";
    assertKey(4096, password, salt, "3d 2e ec 4f e4 1c 84 9b 80 c8 d8 36 62 c0 e4 4a 8b 29 1a 96 4c f2 f0 70 38");

    password = "pass\0word";
    salt = "sa\0lt";
    assertKey(4096, password, salt, "56 fa 6a a7 55 48 09 9d cc 37 d7 f0 34 25 e0 c3");
  }

  private void assertKey(int iterations, String expected) throws UnsupportedEncodingException {
    assertKey(iterations, "password", "salt", expected);
  }

  private void assertKey(int iterations, String password, String salt, String expected) throws UnsupportedEncodingException {
    byte[] expectedBytes = convertHex(expected);
    PBEKeySpec spec = new PBEKeySpec(password.toCharArray(), salt.getBytes("UTF-8"), iterations, expectedBytes.length*8);
    byte[] key = KeyczarPBEReader.pbkdf2(spec);
    assertArrayEquals(expectedBytes, key);
  }

  @Test
  public void testReadEncryptedKey() throws KeyczarException {
    KeyczarPBEReader encrypted = new KeyczarPBEReader(new EncryptedTestKeyReader(), PASSPHRASE);
    String prefix = "{\"aesKeyString\":\"oThFEDqkkLyp80hhh1QFjA\"";
    assertEquals(prefix, encrypted.getKey().substring(0, prefix.length()));

    GenericKeyczar keyczar = new GenericKeyczar(encrypted);
    assertEquals(1, keyczar.getVersions().size());
    assertEquals(1, keyczar.getMetadata().getPrimaryVersion().getVersionNumber());
  }

  @Test
  public void testBadPassword() throws KeyczarException {
    KeyczarPBEReader encrypted = new KeyczarPBEReader(
        new EncryptedTestKeyReader(), "wrong password");
    try {
      encrypted.getKey();
      fail("expected exception");
    } catch (KeyczarException e) {
      assertEquals("Error decrypting PBE key", e.getMessage());
    }
  }

  @Test
  public void testReadUnencryptedKey() throws KeyczarException {
    KeyczarPBEReader reader = new KeyczarPBEReader(new KeyczarFileReader(TEST_DATA + "/aes"), PASSPHRASE);
    try {
      reader.getKey(2);
      fail("expected exception");
    } catch (KeyczarException e) {
      assertEquals(e.getMessage(), "Invalid password-based key metadata");
    }
  }

  @Test
  public void testEncryptKey() throws KeyczarException {
    GenericKeyczar key = generateKey(KeyPurpose.DECRYPT_AND_ENCRYPT, DefaultKeyType.AES);

    KeyczarKey k = key.getKey(key.getMetadata().getPrimaryVersion());
    String unencryptedKey = k.toString();
    String encrypted = KeyczarPBEReader.encryptKey(unencryptedKey, PASSPHRASE);

    PBEKeyczarKey pbeKey = KeyczarPBEReader.parsePBEMetadata(encrypted);
    assertEquals(KeyczarPBEReader.DEFAULT_ITERATION_COUNT, pbeKey.iterationCount);
    assertEquals(KeyczarPBEReader.SALT_BYTES, Base64Coder.decodeWebSafe(pbeKey.salt).length);
    assertEquals(KeyczarPBEReader.PBE_AES_KEY_BYTES, Base64Coder.decodeWebSafe(pbeKey.iv).length);

    // decrypt the key
    KeyczarPBEReader reader = new KeyczarPBEReader(null, PASSPHRASE);
    String out = reader.decryptKey(encrypted);
    assertEquals(unencryptedKey, out);
  }

  /** Generates a new key in memory. */
  private GenericKeyczar generateKey(KeyPurpose purpose, DefaultKeyType type)
      throws KeyczarException {
    KeyMetadata metadata = new KeyMetadata("Key", purpose, type);
    KeyczarReader reader = new ImportedKeyReader(metadata, null);
    GenericKeyczar keyczar = new GenericKeyczar(reader);
    keyczar.addVersion(KeyStatus.PRIMARY);
    return keyczar;
  }
}
